{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 5 - Convolutional Sequence to Sequence Learning\n",
    "\n",
    "In this notebook we'll be implementing the [Convolutional Sequence to Sequence Learning](https://arxiv.org/abs/1705.03122) model. \n",
    "\n",
    "## Introduction\n",
    "\n",
    "## Preparing the Data\n",
    "\n",
    "First, let's import all the required modules and set the random seeds for reproducability."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from torchtext.datasets import TranslationDataset, Multi30k\n",
    "from torchtext.data import Field, BucketIterator\n",
    "\n",
    "import spacy\n",
    "\n",
    "import random\n",
    "import math\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "SEED = 1234\n",
    "\n",
    "random.seed(SEED)\n",
    "torch.manual_seed(SEED)\n",
    "torch.backends.cudnn.deterministic = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we'll load the spaCy models and define the tokenizers for the source and target languages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "spacy_de = spacy.load('de')\n",
    "spacy_en = spacy.load('en')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def tokenize_de(text):\n",
    "    \"\"\"\n",
    "    Tokenizes German text from a string into a list of strings\n",
    "    \"\"\"\n",
    "    return [tok.text for tok in spacy_de.tokenizer(text)]\n",
    "\n",
    "def tokenize_en(text):\n",
    "    \"\"\"\n",
    "    Tokenizes English text from a string into a list of strings\n",
    "    \"\"\"\n",
    "    return [tok.text for tok in spacy_en.tokenizer(text)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we'll set up the `Field`s which decide how the data will be processed. By default RNN models in PyTorch require the sequence to be a tensor of shape **[sequence length, batch size]** so TorchText will, by default, return batches of tensors in the same shape. However in this notebook we are using CNNs which expect the batch dimension to be first. We tell TorchText to have batches be **[batch size, sequence length]** by setting `batch_first = True`.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "SRC = Field(tokenize = tokenize_de, \n",
    "            init_token = '<sos>', \n",
    "            eos_token = '<eos>', \n",
    "            lower = True, \n",
    "            batch_first = True)\n",
    "\n",
    "TRG = Field(tokenize = tokenize_en, \n",
    "            init_token = '<sos>', \n",
    "            eos_token = '<eos>', \n",
    "            lower = True, \n",
    "            batch_first = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, we load our dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data, valid_data, test_data = Multi30k.splits(exts=('.de', '.en'), \n",
    "                                                    fields=(SRC, TRG))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We build our vocabulary as before, by converting any tokens that appear less than 2 times into `<unk>` tokens."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "SRC.build_vocab(train_data, min_freq = 2)\n",
    "TRG.build_vocab(train_data, min_freq = 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The final bit of data preparation is defining the device and then building the iterator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE = 128\n",
    "\n",
    "train_iterator, valid_iterator, test_iterator = BucketIterator.splits(\n",
    "    (train_data, valid_data, test_data), \n",
    "     batch_size = BATCH_SIZE,\n",
    "     device = device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building the Model\n",
    "\n",
    "Next up is building the model. As before, the model is made of an encoder and decoder. The encoder *encodes* the input sentence, in the source language, into a *context vector*. The decoder *decodes* the context vector to produce the output sentence in the target language.\n",
    "\n",
    "### Encoder\n",
    "\n",
    "Previous models in these tutorials had an encoder that compresses an entire input sentence into a single context vector, $z$. The convolutional sequence-to-sequence model is a little different - it gets two context vectors for each token in the input sentence. So, if our input sentence had 6 tokens, we would get 12 context vectors, two for each token. \n",
    "\n",
    "The two context vectors per token are a *conved* vector and a *combined* vector. The conved vector is the result of each token being passed through a few layers - which we will explain shortly. The combined vector comes from the sum of the convolved vector and the embedding of that token.\n",
    "\n",
    "The image below shows the result of an input sentence - *zwei menschen fechten.* - being passed through the encoder.\n",
    "\n",
    "![](assets/convseq2seq1.png)\n",
    "\n",
    "First, the token is passed through a *token embedding layer* - which is standard for neural networks in natural language processing. However, as there are no recurrent connections in this model it has no idea about the order of the tokens within a sequence. To rectify this we have a second embedding layer, the *positional embedding layer*. This is a standard embedding layer where the input is not the token itself but the position of the token within the sequence - starting with the first token, the `<sos>` (start of sequence) token, in position 0.\n",
    "\n",
    "Next, the token and positional embeddings are elementwise summed together to get a vector which contains information about the token and also its position with in the sequence - which we simply call the *embedding vector*. This is followed by a linear layer which transforms the embedding vector into a vector with the required hidden dimension size. \n",
    "\n",
    "The next step is to pass this hidden vector into $N$ *convolutional blocks*. This is where the \"magic\" happens in this model and we will detail the contents of the convolutional blocks shortly. After passing through the convolutional blocks, the vector is then fed through another linear layer to transform it back from the hidden dimension size into the embedding dimension size. This is our *conved* vector - and we have one of these per token in the input sequence. \n",
    "\n",
    "Finally, the conved vector is elementwise summed with the embedding vector via a residual connection to get a *combined* vector for each token. Again, there is a combined vector for each token in the input sequence.\n",
    "\n",
    "### Convolutional Blocks\n",
    "\n",
    "So, how do these convolutional blocks work? The below image shows 2 convolutional blocks, i.e. $N = 2$.\n",
    "\n",
    "![](assets/convseq2seq2.png)\n",
    "\n",
    "First, the input sentence is padded. This is because the convolutional layers will reduce the length of the input sentence and we want the length of the sentence coming into the convolutional blocks to equal the length of it coming into the convolutional blocks. Without padding, the length of the sequence coming out of a convolutional layer will be `filter_size - 1` shorter than the sequence entering the convolutional layer. For example, if we had a filter size of 3, the sequence will be 2 elements shorter. Thus, we pad the sentence with one padding element on each side. We can calculate the amount of padding on each side by simply doing `(filter_size - 1)/2` for odd sized filters - we will not cover even sized filters in this tutorial.\n",
    "\n",
    "These filters are designed so the output hidden dimension of them is twice the input hidden dimension. In computer vision terminology these hidden dimensions are called *channels* - but we will stick to calling them hidden dimensions. Why do we double the size of the hidden dimension leaving the convolutional filter? This is because we are using a special activation function called *gated linear units* (GLU). GLUs have gating mechanisms (similar to LSTMs and GRUs) contained within the activation function and actually half the size of the hidden dimension - whereas usually activation functions keep the hidden dimensions the same size.\n",
    "\n",
    "After passing through the GLU activation the hidden dimension size for each token is the same as it was when it entered the convolutional blocks. It is now elementwise summed with its own vector before it was passed through the convolutional layer. \n",
    "\n",
    "This concludes a single convolutional block. Subsequent blocks take the output of the previous block and perform the same steps. Each block has their own parameters, they are not shared between blocks. The output of the last block goes back to the main encoder - where it is fed through a linear layer to get the conved output and then elementwise summed with the embedding of the token to get the combined output.\n",
    "\n",
    "### Encoder Implementation\n",
    "\n",
    "To keep the implementation simple, we only allow for odd sized kernels. This allows padding to be added equally to both sides of the source sequence.\n",
    "\n",
    "The `scale` variable is used by the authors to \"ensure that the variance throughout the network does not change dramatically\". The performance of the model seems to vary wildly using different seeds if this is not used.\n",
    "\n",
    "The positional embedding is initialized to have a \"vocabulary\" of 100. This means it can handle sequences up to 100 elements long, indexed from 0 to 99. This can be increased if used on a dataset with longer sequences."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Encoder(nn.Module):\n",
    "    def __init__(self, input_dim, emb_dim, hid_dim, n_layers, kernel_size, dropout, device):\n",
    "        super().__init__()\n",
    "        \n",
    "        assert kernel_size % 2 == 1, \"Kernel size must be odd!\"\n",
    "        \n",
    "        self.device = device\n",
    "        \n",
    "        self.scale = torch.sqrt(torch.FloatTensor([0.5])).to(device)\n",
    "        \n",
    "        self.tok_embedding = nn.Embedding(input_dim, emb_dim)\n",
    "        self.pos_embedding = nn.Embedding(100, emb_dim)\n",
    "        \n",
    "        self.emb2hid = nn.Linear(emb_dim, hid_dim)\n",
    "        self.hid2emb = nn.Linear(hid_dim, emb_dim)\n",
    "        \n",
    "        self.convs = nn.ModuleList([nn.Conv1d(in_channels = hid_dim, \n",
    "                                              out_channels = 2 * hid_dim, \n",
    "                                              kernel_size = kernel_size, \n",
    "                                              padding = (kernel_size - 1) // 2)\n",
    "                                    for _ in range(n_layers)])\n",
    "        \n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        \n",
    "    def forward(self, src):\n",
    "        \n",
    "        #src = [batch size, src sent len]\n",
    "        \n",
    "        #create position tensor\n",
    "        pos = torch.arange(0, src.shape[1]).unsqueeze(0).repeat(src.shape[0], 1).to(self.device)\n",
    "        \n",
    "        #pos = [0, 1, 2, 3, ..., src sent len - 1]\n",
    "        \n",
    "        #pos = [batch size, src sent len]\n",
    "        \n",
    "        #embed tokens and positions\n",
    "        tok_embedded = self.tok_embedding(src)\n",
    "        pos_embedded = self.pos_embedding(pos)\n",
    "        \n",
    "        #tok_embedded = pos_embedded = [batch size, src sent len, emb dim]\n",
    "        \n",
    "        #combine embeddings by elementwise summing\n",
    "        embedded = self.dropout(tok_embedded + pos_embedded)\n",
    "        \n",
    "        #embedded = [batch size, src sent len, emb dim]\n",
    "        \n",
    "        #pass embedded through linear layer to convert from emb dim to hid dim\n",
    "        conv_input = self.emb2hid(embedded)\n",
    "        \n",
    "        #conv_input = [batch size, src sent len, hid dim]\n",
    "        \n",
    "        #permute for convolutional layer\n",
    "        conv_input = conv_input.permute(0, 2, 1) \n",
    "        \n",
    "        #conv_input = [batch size, hid dim, src sent len]\n",
    "        \n",
    "        #begin convolutional blocks...\n",
    "        \n",
    "        for i, conv in enumerate(self.convs):\n",
    "        \n",
    "            #pass through convolutional layer\n",
    "            conved = conv(self.dropout(conv_input))\n",
    "\n",
    "            #conved = [batch size, 2 * hid dim, src sent len]\n",
    "\n",
    "            #pass through GLU activation function\n",
    "            conved = F.glu(conved, dim = 1)\n",
    "\n",
    "            #conved = [batch size, hid dim, src sent len]\n",
    "            \n",
    "            #apply residual connection\n",
    "            conved = (conved + conv_input) * self.scale\n",
    "\n",
    "            #conved = [batch size, hid dim, src sent len]\n",
    "            \n",
    "            #set conv_input to conved for next loop iteration\n",
    "            conv_input = conved\n",
    "        \n",
    "        #...end convolutional blocks\n",
    "        \n",
    "        #permute and convert back to emb dim\n",
    "        conved = self.hid2emb(conved.permute(0, 2, 1))\n",
    "        \n",
    "        #conved = [batch size, src sent len, emb dim]\n",
    "        \n",
    "        #elementwise sum output (conved) and input (embedded) to be used for attention\n",
    "        combined = (conved + embedded) * self.scale\n",
    "        \n",
    "        #combined = [batch size, src sent len, emb dim]\n",
    "        \n",
    "        return conved, combined"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Decoder\n",
    "\n",
    "The decoder takes in the actual target sentence and tries to predict it. This model differs from the recurrent neural network models previously detailed in these tutorials as it predicts all tokens within the target sentence in parallel. There is no sequential processing, i.e. no decoding loop. This will be detailed further later on in the tutorials.\n",
    "\n",
    "The decoder is similar to the encoder, with a few changes to both the main model and the convolutional blocks inside the model.\n",
    "\n",
    "![](assets/convseq2seq3.png)\n",
    "\n",
    "First, the embeddings do not have a residual connection that connects after the convolutional blocks and the transformation. Instead the embeddings are fed into the convolutional blocks to be used as residual connections there.\n",
    "\n",
    "Second, to feed the decoder information from the encoder, the encoder conved and combined outputs are used - again, within the convolutional blocks. \n",
    "\n",
    "Finally, the output of the decoder is a linear layer from embedding dimension to output dimension. This is used make a prediction about what the next word in the translation should be.\n",
    "\n",
    "### Decoder Convolutional Blocks\n",
    "\n",
    "Again, these are similar to the convolutional blocks within the encoder, with a few changes.\n",
    "\n",
    "![](assets/convseq2seq4.png)\n",
    "\n",
    "First, the padding. Instead of padding equally on each side to ensure the length of the sentence stays the same throughout, we only pad at the beginning of the sentence. As we are processing all of the targets simultaneously in parallel, and not sequentially, we need a method of only allowing the filter translating token $i$ to only look at tokens before word $i$. If it was allowed to look at token $i+1$ (the token it should be outputting), the neural network will simply learn to output the next word in the sequence, without actually learning how to translate.\n",
    "\n",
    "With padding, the first filter, which is trying to output `two` using `<sos>` and the encoded representations cannot directly see the word `two` at the next position. Without padding, it would be able to see the word `two` and the model would learn to simply copy the word, without making use of the encoded representations and thus not actually learning how to translate.\n",
    "\n",
    "Second, after the GLU activation and before the residual connection, the block calculates and applies attention - using the encoded representations and the embedding of the current word. Note: we only show the connections to the rightmost token, but they are actually connected to all tokens - this was done for clarity. Each token input uses their own, and only their own, embedding for their own attention calculation.\n",
    "\n",
    "The attention is calculated by first using a linear layer to change the hidden dimension to the same size as the embedding dimension. Then the embedding summed via a residual connection. This combination then has the standard attention calculation applied by finding how much it \"matches\" with the *encoded conved* and then this is applied by getting a weighted sum over the *encoded combined*. This is then projected back up to the hidden dimenson size and a residual connection to the initial input to the attention layer is applied.\n",
    "\n",
    "Why do they calculate attention first with the encoded conved and then use it to calculate the weighted sum over the encoded combined? They argue that the encoded conved is good for getting a larger context over the encoded sequence, whereas the encoded combined has more information about the specific token and is thus therefore more useful for makng a prediction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Decoder(nn.Module):\n",
    "    def __init__(self, output_dim, emb_dim, hid_dim, n_layers, kernel_size, dropout, pad_idx, device):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.kernel_size = kernel_size\n",
    "        self.pad_idx = pad_idx\n",
    "        self.device = device\n",
    "        \n",
    "        self.scale = torch.sqrt(torch.FloatTensor([0.5])).to(device)\n",
    "        \n",
    "        self.tok_embedding = nn.Embedding(output_dim, emb_dim)\n",
    "        self.pos_embedding = nn.Embedding(100, emb_dim)\n",
    "        \n",
    "        self.emb2hid = nn.Linear(emb_dim, hid_dim)\n",
    "        self.hid2emb = nn.Linear(hid_dim, emb_dim)\n",
    "        \n",
    "        self.attn_hid2emb = nn.Linear(hid_dim, emb_dim)\n",
    "        self.attn_emb2hid = nn.Linear(emb_dim, hid_dim)\n",
    "        \n",
    "        self.out = nn.Linear(emb_dim, output_dim)\n",
    "        \n",
    "        self.convs = nn.ModuleList([nn.Conv1d(in_channels = hid_dim, \n",
    "                                              out_channels = 2 * hid_dim, \n",
    "                                              kernel_size = kernel_size)\n",
    "                                    for _ in range(n_layers)])\n",
    "        \n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "      \n",
    "    def calculate_attention(self, embedded, conved, encoder_conved, encoder_combined):\n",
    "        \n",
    "        #embedded = [batch size, trg sent len, emb dim]\n",
    "        #conved = [batch size, hid dim, trg sent len]\n",
    "        #encoder_conved = encoder_combined = [batch size, src sent len, emb dim]\n",
    "        \n",
    "        #permute and convert back to emb dim\n",
    "        conved_emb = self.attn_hid2emb(conved.permute(0, 2, 1))\n",
    "        \n",
    "        #conved_emb = [batch size, trg sent len, emb dim]\n",
    "        \n",
    "        combined = (conved_emb + embedded) * self.scale\n",
    "        \n",
    "        #combined = [batch size, trg sent len, emb dim]\n",
    "                \n",
    "        energy = torch.matmul(combined, encoder_conved.permute(0, 2, 1))\n",
    "        \n",
    "        #energy = [batch size, trg sent len, src sent len]\n",
    "        \n",
    "        attention = F.softmax(energy, dim=2)\n",
    "        \n",
    "        #attention = [batch size, trg sent len, src sent len]\n",
    "            \n",
    "        attended_encoding = torch.matmul(attention, encoder_combined)\n",
    "        \n",
    "        #attended_encoding = [batch size, trg sent len, emd dim]\n",
    "        \n",
    "        #convert from emb dim -> hid dim\n",
    "        attended_encoding = self.attn_emb2hid(attended_encoding)\n",
    "        \n",
    "        #attended_encoding = [batch size, trg sent len, hid dim]\n",
    "        \n",
    "        #apply residual connection\n",
    "        attended_combined = (conved + attended_encoding.permute(0, 2, 1)) * self.scale\n",
    "        \n",
    "        #attended_combined = [batch size, hid dim, trg sent len]\n",
    "        \n",
    "        return attention, attended_combined\n",
    "        \n",
    "    def forward(self, trg, encoder_conved, encoder_combined):\n",
    "        \n",
    "        #trg = [batch size, trg sent len]\n",
    "        #encoder_conved = encoder_combined = [batch size, src sent len, emb dim]\n",
    "                \n",
    "        #create position tensor\n",
    "        pos = torch.arange(0, trg.shape[1]).unsqueeze(0).repeat(trg.shape[0], 1).to(device)\n",
    "        \n",
    "        #pos = [batch size, trg sent len]\n",
    "        \n",
    "        #embed tokens and positions\n",
    "        tok_embedded = self.tok_embedding(trg)\n",
    "        pos_embedded = self.pos_embedding(pos)\n",
    "        \n",
    "        #tok_embedded = [batch size, trg sent len, emb dim]\n",
    "        #pos_embedded = [batch size, trg sent len, emb dim]\n",
    "        \n",
    "        #combine embeddings by elementwise summing\n",
    "        embedded = self.dropout(tok_embedded + pos_embedded)\n",
    "        \n",
    "        #embedded = [batch size, trg sent len, emb dim]\n",
    "        \n",
    "        #pass embedded through linear layer to go through emb dim -> hid dim\n",
    "        conv_input = self.emb2hid(embedded)\n",
    "        \n",
    "        #conv_input = [batch size, trg sent len, hid dim]\n",
    "        \n",
    "        #permute for convolutional layer\n",
    "        conv_input = conv_input.permute(0, 2, 1) \n",
    "        \n",
    "        #conv_input = [batch size, hid dim, trg sent len]\n",
    "        \n",
    "        for i, conv in enumerate(self.convs):\n",
    "        \n",
    "            #apply dropout\n",
    "            conv_input = self.dropout(conv_input)\n",
    "        \n",
    "            #need to pad so decoder can't \"cheat\"\n",
    "            padding = torch.zeros(conv_input.shape[0], \n",
    "                                  conv_input.shape[1], \n",
    "                                  self.kernel_size - 1).fill_(self.pad_idx).to(self.device)\n",
    "                \n",
    "            padded_conv_input = torch.cat((padding, conv_input), dim = 2)\n",
    "        \n",
    "            #padded_conv_input = [batch size, hid dim, trg sent len + kernel size - 1]\n",
    "        \n",
    "            #pass through convolutional layer\n",
    "            conved = conv(padded_conv_input)\n",
    "\n",
    "            #conved = [batch size, 2 * hid dim, trg sent len]\n",
    "            \n",
    "            #pass through GLU activation function\n",
    "            conved = F.glu(conved, dim = 1)\n",
    "\n",
    "            #conved = [batch size, hid dim, trg sent len]\n",
    "            \n",
    "            #calculate attention\n",
    "            attention, conved = self.calculate_attention(embedded, conved, encoder_conved, encoder_combined)\n",
    "            \n",
    "            #attention = [batch size, trg sent len, src sent len]\n",
    "            \n",
    "            #apply residual connection\n",
    "            conved = (conved + conv_input) * self.scale\n",
    "            \n",
    "            #conved = [batch size, hid dim, trg sent len]\n",
    "            \n",
    "            #set conv_input to conved for next loop iteration\n",
    "            conv_input = conved\n",
    "            \n",
    "        conved = self.hid2emb(conved.permute(0, 2, 1))\n",
    "         \n",
    "        #conved = [batch size, trg sent len, emb dim]\n",
    "            \n",
    "        output = self.out(self.dropout(conved))\n",
    "        \n",
    "        #output = [batch size, trg sent len, output dim]\n",
    "            \n",
    "        return output, attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Seq2Seq(nn.Module):\n",
    "    def __init__(self, encoder, decoder, device):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "        self.device = device\n",
    "        \n",
    "    def forward(self, src, trg):\n",
    "        \n",
    "        #src = [batch size, src sent len]\n",
    "        #trg = [batch size, trg sent len]\n",
    "           \n",
    "        #calculate z^u (encoder_conved) and e (encoder_combined)\n",
    "        #encoder_conved is output from final encoder conv. block\n",
    "        #encoder_combined is encoder_conved plus (elementwise) src embedding plus positional embeddings \n",
    "        encoder_conved, encoder_combined = self.encoder(src)\n",
    "            \n",
    "        #encoder_conved = [batch size, src sent len, emb dim]\n",
    "        #encoder_combined = [batch size, src sent len, emb dim]\n",
    "        \n",
    "        #calculate predictions of next words\n",
    "        #output is a batch of predictions for each word in the trg sentence\n",
    "        #attention a batch of attention scores across the src sentence for each word in the trg sentence\n",
    "        output, attention = self.decoder(trg, encoder_conved, encoder_combined)\n",
    "        \n",
    "        #output = [batch size, trg sent len, output dim]\n",
    "        #attention = [batch size, trg sent len, src sent len]\n",
    "        \n",
    "        return output, attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "INPUT_DIM = len(SRC.vocab)\n",
    "OUTPUT_DIM = len(TRG.vocab)\n",
    "EMB_DIM = 256\n",
    "HID_DIM = 512\n",
    "ENC_LAYERS = 10\n",
    "DEC_LAYERS = 10\n",
    "ENC_KERNEL_SIZE = 3\n",
    "DEC_KERNEL_SIZE = 3\n",
    "ENC_DROPOUT = 0.25\n",
    "DEC_DROPOUT = 0.25\n",
    "PAD_IDX = TRG.vocab.stoi['<pad>']\n",
    "    \n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "    \n",
    "enc = Encoder(INPUT_DIM, EMB_DIM, HID_DIM, ENC_LAYERS, ENC_KERNEL_SIZE, ENC_DROPOUT, device)\n",
    "dec = Decoder(OUTPUT_DIM, EMB_DIM, HID_DIM, DEC_LAYERS, DEC_KERNEL_SIZE, DEC_DROPOUT, PAD_IDX, device)\n",
    "\n",
    "model = Seq2Seq(enc, dec, device).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The model has 37,351,685 trainable parameters\n"
     ]
    }
   ],
   "source": [
    "def count_parameters(model):\n",
    "    return sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
    "\n",
    "print(f'The model has {count_parameters(model):,} trainable parameters')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = optim.Adam(model.parameters())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "criterion = nn.CrossEntropyLoss(ignore_index = PAD_IDX)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(model, iterator, optimizer, criterion, clip):\n",
    "    \n",
    "    model.train()\n",
    "    \n",
    "    epoch_loss = 0\n",
    "    \n",
    "    for i, batch in enumerate(iterator):\n",
    "        \n",
    "        src = batch.src\n",
    "        trg = batch.trg\n",
    "        \n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        output, _ = model(src, trg[:,:-1])\n",
    "        \n",
    "        #output = [batch size, trg sent len - 1, output dim]\n",
    "        #trg = [batch size, trg sent len]\n",
    "        \n",
    "        output = output.contiguous().view(-1, output.shape[-1])\n",
    "        trg = trg[:,1:].contiguous().view(-1)\n",
    "        \n",
    "        #output = [batch size * trg sent len - 1, output dim]\n",
    "        #trg = [batch size * trg sent len - 1]\n",
    "        \n",
    "        loss = criterion(output, trg)\n",
    "        \n",
    "        loss.backward()\n",
    "        \n",
    "        torch.nn.utils.clip_grad_norm_(model.parameters(), clip)\n",
    "        \n",
    "        optimizer.step()\n",
    "        \n",
    "        epoch_loss += loss.item()\n",
    "        \n",
    "    return epoch_loss / len(iterator)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(model, iterator, criterion):\n",
    "    \n",
    "    model.eval()\n",
    "    \n",
    "    epoch_loss = 0\n",
    "    \n",
    "    with torch.no_grad():\n",
    "    \n",
    "        for i, batch in enumerate(iterator):\n",
    "\n",
    "            src = batch.src\n",
    "            trg = batch.trg\n",
    "\n",
    "            output, _ = model(src, trg[:,:-1])\n",
    "        \n",
    "            #output = [batch size, trg sent len - 1, output dim]\n",
    "            #trg = [batch size, trg sent len]\n",
    "\n",
    "            output = output.contiguous().view(-1, output.shape[-1])\n",
    "            trg = trg[:,1:].contiguous().view(-1)\n",
    "\n",
    "            #output = [batch size * trg sent len - 1, output dim]\n",
    "            #trg = [batch size * trg sent len - 1]\n",
    "            \n",
    "            loss = criterion(output, trg)\n",
    "\n",
    "            epoch_loss += loss.item()\n",
    "        \n",
    "    return epoch_loss / len(iterator)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "def epoch_time(start_time, end_time):\n",
    "    elapsed_time = end_time - start_time\n",
    "    elapsed_mins = int(elapsed_time / 60)\n",
    "    elapsed_secs = int(elapsed_time - (elapsed_mins * 60))\n",
    "    return elapsed_mins, elapsed_secs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 01 | Time: 0m 30s\n",
      "\tTrain Loss: 4.202 | Train PPL:  66.837\n",
      "\t Val. Loss: 2.907 |  Val. PPL:  18.306\n",
      "Epoch: 02 | Time: 0m 30s\n",
      "\tTrain Loss: 3.018 | Train PPL:  20.446\n",
      "\t Val. Loss: 2.378 |  Val. PPL:  10.779\n",
      "Epoch: 03 | Time: 0m 30s\n",
      "\tTrain Loss: 2.598 | Train PPL:  13.438\n",
      "\t Val. Loss: 2.121 |  Val. PPL:   8.340\n",
      "Epoch: 04 | Time: 0m 30s\n",
      "\tTrain Loss: 2.360 | Train PPL:  10.591\n",
      "\t Val. Loss: 1.976 |  Val. PPL:   7.210\n",
      "Epoch: 05 | Time: 0m 30s\n",
      "\tTrain Loss: 2.208 | Train PPL:   9.102\n",
      "\t Val. Loss: 1.899 |  Val. PPL:   6.678\n",
      "Epoch: 06 | Time: 0m 30s\n",
      "\tTrain Loss: 2.095 | Train PPL:   8.124\n",
      "\t Val. Loss: 1.845 |  Val. PPL:   6.328\n",
      "Epoch: 07 | Time: 0m 30s\n",
      "\tTrain Loss: 2.008 | Train PPL:   7.450\n",
      "\t Val. Loss: 1.801 |  Val. PPL:   6.059\n",
      "Epoch: 08 | Time: 0m 30s\n",
      "\tTrain Loss: 1.933 | Train PPL:   6.909\n",
      "\t Val. Loss: 1.776 |  Val. PPL:   5.908\n",
      "Epoch: 09 | Time: 0m 30s\n",
      "\tTrain Loss: 1.872 | Train PPL:   6.499\n",
      "\t Val. Loss: 1.748 |  Val. PPL:   5.743\n",
      "Epoch: 10 | Time: 0m 30s\n",
      "\tTrain Loss: 1.819 | Train PPL:   6.166\n",
      "\t Val. Loss: 1.731 |  Val. PPL:   5.649\n"
     ]
    }
   ],
   "source": [
    "N_EPOCHS = 10\n",
    "CLIP = 1\n",
    "\n",
    "best_valid_loss = float('inf')\n",
    "\n",
    "for epoch in range(N_EPOCHS):\n",
    "    \n",
    "    start_time = time.time()\n",
    "    \n",
    "    train_loss = train(model, train_iterator, optimizer, criterion, CLIP)\n",
    "    valid_loss = evaluate(model, valid_iterator, criterion)\n",
    "    \n",
    "    end_time = time.time()\n",
    "    \n",
    "    epoch_mins, epoch_secs = epoch_time(start_time, end_time)\n",
    "    \n",
    "    if valid_loss < best_valid_loss:\n",
    "        best_valid_loss = valid_loss\n",
    "        torch.save(model.state_dict(), 'tut5-model.pt')\n",
    "    \n",
    "    print(f'Epoch: {epoch+1:02} | Time: {epoch_mins}m {epoch_secs}s')\n",
    "    print(f'\\tTrain Loss: {train_loss:.3f} | Train PPL: {math.exp(train_loss):7.3f}')\n",
    "    print(f'\\t Val. Loss: {valid_loss:.3f} |  Val. PPL: {math.exp(valid_loss):7.3f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "| Test Loss: 1.800 | Test PPL:   6.047 |\n"
     ]
    }
   ],
   "source": [
    "model.load_state_dict(torch.load('tut5-model.pt'))\n",
    "\n",
    "test_loss = evaluate(model, test_iterator, criterion)\n",
    "\n",
    "print(f'| Test Loss: {test_loss:.3f} | Test PPL: {math.exp(test_loss):7.3f} |')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
